from common import cut_text
from sklearn.metrics import accuracy_score, recall_score
import pandas as pd
import pickle
import seaborn as sns


# 从文件读取训练好的模型 进行预测
def use_model_predict(sname, model_file_name, model_dir, path, header):
    # 从文件读取TF-IDF模型
    with open("../model/TF-IDF/%s_TF_IDF_MODEL.pkl" % sname, "rb") as f:
        tf_idf_model = pickle.load(f)

    # 导入停用词表
    with open("../data/stop_words.txt", 'r', encoding="utf-8-sig") as f:
        line = f.readline()
        stopwords = line.split(",")

    # 从文件读取模型
    with open("../model/%s/%s" % (model_dir, model_file_name), "rb") as f:
        mnb_model = pickle.load(f)

    # 读取数据
    df = pd.read_csv(path, names=header)
    # 随机取10%数据作为测试集
    # dfSample = df.sample(frac=1)
    dfSample = df

    # 分词
    dfSampleCut = dfSample['review'].map(lambda x: cut_text(x, stopwords))

    # 正确结果
    true_y = dfSample['label']
    # 分词转TF-IDF后 使用模型进行预测得到预测结果
    pred_y = mnb_model.predict(tf_idf_model.transform(dfSampleCut))

    return true_y, pred_y


# 根据预测结果 计算指标
def cal_metrics(sname, model_file_name, model_dir, path, header, repeat=1):
    sum_accuracy_score = 0
    sum_recall_score = 0
    # 进行50次随机取样预测，计算准确率平均值
    for i in range(repeat):
        true_y, pred_y = use_model_predict(sname=sname, model_file_name=model_file_name, model_dir=model_dir, path=path,
                                           header=header)
        print(true_y.value_counts())
        print(pd.Series(pred_y).value_counts())
        # 计算准确率
        sum_accuracy_score += accuracy_score(true_y, pred_y)
        # 计算召回率
        sum_recall_score += recall_score(true_y, pred_y)

    # 平均准确率
    avg_accuracy_score = sum_accuracy_score / repeat
    # 平均召回率
    avg_recall_score = sum_recall_score / repeat
    return avg_accuracy_score, avg_recall_score


# 根据数据类型获取数据路径
def getDataPathAndHeader(sname):
    header = ['label', 'review']
    if sname == 'weibo':
        return "../data/weibo_senti_100k.csv", header
    elif sname == 'waimai':
        return "../data/waimai_10k.csv", header
    elif sname == 'shopping':
        header = ['cat', 'label', 'review']
        return "../data/online_shopping_10_cats.csv", header


# 根据模型名成、数据名称获取模型路径
def getModelFileName(model_name, sname):
    if model_name == 'Bayes':
        return sname + "_MNB_MODEL.pkl"
    elif model_name == 'AdaBoostBayes':
        return sname + "_ADABOOST_MNB_MODEL.pkl"
    elif model_name == 'DecisionTree':
        return sname + "_DT_MODEL.pkl"
    elif model_name == 'RandomForest':
        return sname + "_RF_MODEL.pkl"
    elif model_name == 'RGF':
        return sname + "_RGF_MODEL.pkl"


# 计算每一种模型的指标
def CalEveryModelMetrixs():
    # 三类数据
    snames = ['weibo', 'waimai', 'shopping']
    # 五大模型dir
    model_names = ['Bayes', 'AdaBoostBayes', 'DecisionTree', 'RandomForest', 'RGF']
    # 保存结果
    data_dir = {}
    # 遍历三类数据集
    for sname in snames:
        model_metrix_dir = {}
        for model_name in model_names:
            model_file_name = getModelFileName(model_name, sname)
            path, header = getDataPathAndHeader(sname)
            accuracy_score, recall_score = cal_metrics(sname=sname, model_file_name=model_file_name,
                                                       model_dir=model_name,
                                                       path=path, header=header)
            model_metrix_dir[model_name] = [accuracy_score, recall_score]
        data_dir[sname] = model_metrix_dir
        print(model_name + "model predict ok")
    return data_dir


# 将结果转换成适合绘图的字典结构
def convert_to_draw(result_dict):
    draw_dict = {}
    for dataType in result_dict:
        tmp_draw_dict = {}
        dataType_list = []
        metrix_list = []
        metrix_name = []
        model_list = []
        for model in result_dict[dataType]:
            metrix_list.append(result_dict[dataType][model][0])
            metrix_name.append('accuracy')
            dataType_list.append(dataType)
            model_list.append(model)
            metrix_list.append(result_dict[dataType][model][1])
            metrix_name.append('recall')
            dataType_list.append(dataType)
            model_list.append(model)
        tmp_draw_dict['metrix_value'] = metrix_list
        tmp_draw_dict['dataType'] = dataType_list
        tmp_draw_dict['metrix_name'] = metrix_name
        tmp_draw_dict['model'] = model_list
        draw_dict[dataType] = tmp_draw_dict
    return draw_dict


# 进行绘图
def draw():
    result_dict = CalEveryModelMetrixs()
    draw_dict = convert_to_draw(result_dict)
    for dataType in draw_dict:
        sns.lineplot(x='model', y='metrix_value', hue='metrix_name', data=pd.DataFrame(draw_dict[dataType]))


def useTrainPredictData():
    # 三类数据
    snames = ['weibo', 'waimai', 'shopping']
    # 五大模型dir
    model_names = ['Bayes', 'AdaBoostBayes', 'DecisionTree', 'RandomForest', 'RGF']
    # 保存结果
    data_dir = {}
    # 遍历三类数据集
    for sname in snames:
        model_metrix_dir = {}
        for model_name in model_names:
            model_file_name = getModelFileName(model_name, 'train')
            path, header = getDataPathAndHeader(sname)
            accuracy_score, recall_score = cal_metrics(sname='train', model_file_name=model_file_name,
                                                       model_dir=model_name,
                                                       path=path, header=header)
            model_metrix_dir[model_name] = [accuracy_score, recall_score]
        data_dir[sname] = model_metrix_dir
    return data_dir


if __name__ == '__main__':
    # CalEveryModelMetrixs()
    print(useTrainPredictData())
